# -*- coding: utf-8 -*-

import torch
import segmentation_models_pytorch as smp
from CC_CCI_dataset import LabeledDataset, visualize
from losses import CrossEntropyLoss
from metrics import Accuracy, IoU


dataset = LabeledDataset("D:\data\ct_lesion_seg\image", "D:\data\ct_lesion_seg\mask")
trainset, validset = torch.utils.data.random_split(dataset, [int(len(dataset)*0.7), len(dataset)-int(len(dataset)*0.7)])
train_loader = torch.utils.data.DataLoader(trainset, batch_size=8, shuffle=True)
valid_loader = torch.utils.data.DataLoader(validset, batch_size=1, shuffle=True)

# med_dataset = KaggleSegDataset("D:\data\covid19-ct-scans\scans", "D:\data\covid19-ct-scans\masks")
# med_trainset, med_validset = torch.utils.data.random_split(med_dataset, [int(len(med_dataset)*0.9), len(med_dataset)-int(len(med_dataset)*0.9)])
# med_train_loader = torch.utils.data.DataLoader(med_trainset, batch_size=8, shuffle=True)
# med_valid_loader = torch.utils.data.DataLoader(med_validset, batch_size=1, shuffle=True)

# med1_dataset = MedSegDataset("D:\data\covid-segmentation\images_medseg.npy", "D:\data\covid-segmentation\masks_medseg.npy")
# med1_trainset, med1_validset = torch.utils.data.random_split(med1_dataset, [int(len(med1_dataset)*0.9), len(med1_dataset)-int(len(med1_dataset)*0.9)])
# med1_train_loader = torch.utils.data.DataLoader(med1_trainset, batch_size=8, shuffle=True)
# med1_valid_loader = torch.utils.data.DataLoader(med1_validset, batch_size=1, shuffle=True)

model = smp.DeepLabV3Plus(encoder_name="mobilenet_v2", encoder_weights=None, in_channels=1, classes=3, activation='softmax2d')

loss = CrossEntropyLoss()
optimizer = torch.optim.Adam(params=model.parameters())
metrics = [
    Accuracy(),
    IoU()
]

train_epoch = smp.utils.train.TrainEpoch(
    model,
    loss=loss,
    metrics=metrics,
    optimizer=optimizer,
    device='cuda',
    verbose=True,
)

valid_epoch = smp.utils.train.ValidEpoch(
    model,
    loss=loss,
    metrics=metrics,
    device='cuda',
    verbose=True,
)

for i in range(0, 10):
    print('\nEpoch: {}'.format(i))
    train_logs = train_epoch.run(train_loader)
    valid_logs = valid_epoch.run(valid_loader)
    # med_train_logs = train_epoch.run(med_train_loader)
    # med_valid_logs = valid_epoch.run(med_valid_loader)
    # med1_train_logs = train_epoch.run(med1_train_loader)
    # med1_valid_logs = valid_epoch.run(med1_valid_loader)
    

def sample_test(loader):
    for i, item in enumerate(loader):
        if i >= 1:
            break
        x, y = item
        
        pred = model(x.cuda())
        pred = torch.argmax(pred.cpu(), dim=1)
        visualize(pred, y.cpu()[0])
        
    return pred[0], y.cpu()[0]